import cv2, taichi as ti, numpy as np

from .common import Vec2, Vec3, mix
from ..config import GAMMA

__all__ = [
    "getImage",
]


def imageRead(filepath: str) -> np.ndarray:
    """读取图片文件"""
    if filepath.endswith(".hdr"):
        img = cv2.imread(filepath, flags=cv2.IMREAD_ANYDEPTH)
    if filepath.endswith(".tif"):
        img = cv2.imread(filepath, flags=cv2.IMREAD_UNCHANGED).astype(np.float32) / 255
    else:
        img = ti.tools.imread(filepath).astype(np.float32) / 255
        img = np.transpose(img, (1, 0, 2))
    return img


def getImage(filepath: str):
    """由图片制作材质贴图"""
    img = imageRead(filepath)

    class Image:
        data = Vec3.field(shape=img.shape[:2])
        data.from_numpy(img)

        @ti.func
        def sample(self, uv: Vec2) -> Vec3:
            uv *= Vec2(Image.data.shape[:2]) - 1
            index = int(uv)
            index_ = index + 1
            dp = uv - index
            y0 = mix(Image.data[index.x, index.y], Image.data[index_.x, index.y], dp.x)
            y1 = mix(
                Image.data[index.x, index_.y], Image.data[index_.x, index_.y], dp.x
            )
            return mix(y0, y1, dp.y)

    return Image()
